#!/usr/bin/env python3
import rospy
import numpy as np
import open3d as o3d
from sensor_msgs.msg import PointCloud2
import sensor_msgs.point_cloud2 as pc2

class RealtimePointCloudVisualizer:
    def __init__(self):
        # 初始化 ROS 节点
        rospy.init_node('pointcloud_visualizer', anonymous=True)
        
        # 参数设置
        self.topic_name = rospy.get_param('~topic_name', '/camera/depth_registered/points')
        self.voxel_size = rospy.get_param('~voxel_size', 0.01)  # 降采样体视素大小（米）
        self.update_rate = rospy.get_param('~update_rate', 10.0)  # 更新频率（Hz）
        self.z_min = rospy.get_param('~z_min', 0.1)  # 最小深度（米）
        self.z_max = rospy.get_param('~z_max', 5.0)  # 最大深度（米）
        self.point_size = rospy.get_param('~point_size', 2.0)  # 点渲染大小
        self.window_width = rospy.get_param('~window_width', 1280)  # 窗口宽度
        self.window_height = rospy.get_param('~window_height', 720)  # 窗口宽度
        
        # 初始化 Open3D 可视化器
        self.vis = o3d.visualization.Visualizer()
        self.vis.create_window(window_name="Realtime Point Cloud", width=self.window_width, height=self.window_height)
        self.pcd = o3d.geometry.PointCloud()
        self.is_first_frame = True
        self.last_update = 0.0
        
        # 设置渲染选项
        render_option = self.vis.get_render_option()
        render_option.point_size = self.point_size
        
        # 订阅点云话题
        self.cloud_sub = rospy.Subscriber(
            self.topic_name,
            PointCloud2,
            self.cloud_callback,
            queue_size=1
        )
        
        rospy.loginfo(f"订阅点云话题: {self.topic_name}")
        rospy.loginfo(f"体视素降采样大小: {self.voxel_size} 米")
        rospy.loginfo(f"更新频率: {self.update_rate} Hz")
        rospy.loginfo(f"深度范围: [{self.z_min}, {self.z_max}] 米")
        rospy.loginfo(f"点渲染大小: {self.point_size}")
        rospy.loginfo(f"窗口分辨率: {self.window_width}x{self.window_height}")
        rospy.loginfo("等待点云数据...")

    def cloud_to_o3d(self, msg):
        """将 ROS PointCloud2 转换为 Open3D 点云"""
        # 打印字段信息
        rospy.loginfo(f"点云字段: {str(msg.fields)}")
        
        # 提取 XYZ 数据
        points = np.array(list(pc2.read_points(msg, skip_nans=True, field_names=("x", "y", "z"))))
        if len(points) == 0:
            rospy.logwarn("收到空点云，跳过")
            return None
        
        # 过滤无效点
        valid_mask = (points[:, 2] > self.z_min) & (points[:, 2] < self.z_max) & \
                     (~np.isnan(points[:, 0])) & (~np.isnan(points[:, 1])) & (~np.isnan(points[:, 2]))
        points = points[valid_mask]
        
        if len(points) == 0:
            rospy.logwarn("所有点无效，跳过")
            return None
        
        # 打印点云坐标范围
        rospy.loginfo(f"点云坐标范围: x=[{points[:, 0].min():.2f}, {points[:, 0].max():.2f}], "
                      f"y=[{points[:, 1].min():.2f}, {points[:, 1].max():.2f}], "
                      f"z=[{points[:, 2].min():.2f}, {points[:, 2].max():.2f}]")
        
        # 创建 Open3D 点云
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(points)
        
        # 处理颜色数据（rgb 为 float32）
        if 'rgb' in str(msg.fields):
            try:
                colors = np.array(list(pc2.read_points(msg, skip_nans=True, field_names=("rgb",))))
                colors = colors.flatten()[valid_mask]
                if len(colors) == len(points):
                    # 假设 rgb 是 float32，值在 0-255 范围
                    colors = colors.astype(np.float32)
                    rgb = np.zeros((len(colors), 3))
                    rgb[:, 0] = (colors / 256.0) & 0xFF  # Red
                    rgb[:, 1] = (colors / 256.0) & 0xFF  # Green
                    rgb[:, 2] = (colors / 256.0) & 0xFF  # Blue
                    pcd.colors = o3d.utility.Vector3dVector(rgb.astype(np.float64) / 255.0)
                    rospy.loginfo(f"处理颜色数据，点数: {len(colors)}")
                else:
                    rospy.logwarn("颜色数据和点云数量不匹配，跳过颜色")
            except Exception as e:
                rospy.logwarn(f"颜色处理失败: {str(e)}，跳过颜色")
        
        # 体视素降采样
        if self.voxel_size > 0:
            pcd = pcd.voxel_down_sample(voxel_size=self.voxel_size)
            rospy.loginfo(f"降采样后点云点数: {len(pcd.points)}")
        
        rospy.loginfo(f"有效点云点数: {len(pcd.points)}")
        return pcd

    def cloud_callback(self, msg):
        """点云回调函数 - 更新并显示点云"""
        # 控制更新频率
        current_time = rospy.get_time()
        if current_time - self.last_update < 1.0 / self.update_rate:
            return
        self.last_update = current_time
        
        try:
            pcd = self.cloud_to_o3d(msg)
            if pcd is None:
                return
                
            # 更新点云
            self.pcd.points = pcd.points
            if pcd.has_colors():
                self.pcd.colors = pcd.colors
            
            # 初始化或更新可视化
            if self.is_first_frame:
                self.vis.add_geometry(self.pcd)
                # 设置固定视角
                ctr = self.vis.get_view_control()
                ctr.set_front([0, 0, -1])
                ctr.set_up([0, -1, 0])
                ctr.set_zoom(0.3)  # 手动设置缩放，适合大多数场景
                self.is_first_frame = False
            else:
                self.vis.update_geometry(self.pcd)
            
            # 更新渲染
            self.vis.poll_events()
            self.vis.update_renderer()
            
        except Exception as e:
            rospy.logerr(f"点云处理失败: {str(e)}")

    def run(self):
        """主循环"""
        try:
            rospy.spin()
        except rospy.ROSInterruptException:
            rospy.loginfo("节点被中断")
        finally:
            self.vis.destroy_window()

if __name__ == "__main__":
    try:
        visualizer = RealtimePointCloudVisualizer()
        visualizer.run()
    except Exception as e:
        rospy.logerr(f"初始化失败: {str(e)}")
